import datetime
from pathlib import Path

import torch
import torchvision
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F
from common import *
# 2d 卷积操作函数
from torch import nn
from torchvision import transforms


# 3x3 convolution
def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3,
                     stride=stride, padding=1, bias=False)


# 残差模块 Residual block
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        # 定义三层卷积
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)  # BatchNorm2d最常用于卷积网络中(防止梯度消失或爆炸)，设置的参数就是卷积的输出通道数
        self.relu = nn.ReLU(inplace=True)

        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)  # BatchNorm2d最常用于卷积网络中(防止梯度消失或爆炸)，设置的参数就是卷积的输出通道数
        self.downsample = downsample

    def forward(self, X):
        residual = X
        out = self.conv1(X)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample:
            residual = self.downsample(X)
        out += residual
        out = self.relu(out)
        return out


class ResidualNet(nn.Module):
    def __init__(self, block, layers, num_classes=3):
        super(ResidualNet, self).__init__()
        self.in_channels = 64
        self.conv = conv3x3(3, 64)
        self.bn = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self.make_layer(block, 64, layers[0])
        self.layer2 = self.make_layer(block, 128, layers[1], 2)
        self.layer3 = self.make_layer(block, 256, layers[2], 2)
        self.layer4 = self.make_layer(block, 512, layers[3], 2)
        self.avg_poll = nn.AvgPool2d(4, ceil_mode=False)
        self.fc = nn.Linear(512, num_classes)

    def make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if (stride != 1) or (self.in_channels != out_channels):
            print(f"self.in_channels={self.in_channels}, out_channels={out_channels}")
            downsample = nn.Sequential(
                conv3x3(self.in_channels, out_channels, stride=stride),
                nn.BatchNorm2d(out_channels)
            )
        layers = []
        # 残差直接映射部分
        layers.append(block(self.in_channels, out_channels, stride=stride, downsample=downsample))
        self.in_channels = out_channels
        for i in range(1, blocks):
            layers.append(block(out_channels, out_channels))
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.avg_poll(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


# 训练函数
def train_epoch(net, data_loader, device, is_train):
    if is_train:
        net.train()  # 指定当前为训练模式
    else:
        net.eval()  # 指定当前为测试模式
    train_batch_num = len(data_loader)  # 记录共有多少个batch
    total_loss = 0  # 记录Loss
    correct = 0  # 记录共有多少个样本被正确分类
    sample_num = 0  # 记录样本总数

    for batch_idx, (data, target) in enumerate(data_loader):
        data = data.to(device).float()
        target = target.to(device).long()

        # 前向传播
        output = net(data)
        loss = criterion(output, target)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        prediction = torch.argmax(output, 1)
        correct += (prediction == target).sum().item()
        sample_num += len(prediction)
    loss = total_loss / train_batch_num
    acc = correct / sample_num
    return loss, acc


# 用于更新参数组中的学习率
def update_lr(optimizer, lr):
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


data_dir = "/Users/ryantong/Documents/bjtu/实验三数据集/residual"
# data_dir = "/Users/ryantong/Documents/bjtu/实验三数据集/车辆分类数据集"
ratio = 0.7
i_size = 32
num_classes = 3
epochs = 60
lr = 0.0005
batch_size = 32
transform = transforms.Compose(
    [transforms.Resize((i_size, i_size), Image.ANTIALIAS)
        , transforms.ToTensor()
        , transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ]
)
all_datas = torchvision.datasets.ImageFolder(data_dir, transform=transform)
classes = all_datas.classes
class_to_idx = all_datas.class_to_idx
sum_len = all_datas.__len__()
# 划分数据集
train_size = int(0.7 * len(all_datas))
test_size = len(all_datas) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(all_datas, [train_size, test_size])
train_dataset = train_dataset.dataset  # 这行很重要
test_dataset = test_dataset.dataset  # 这行很重要
# for images, labels in train_dataset.imgs:
#     print(images)
#     print(labels)

device = torch.device("cpu")
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
layers = [2, 2, 2, 2]
net = ResidualNet(ResidualBlock, layers, num_classes).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
ExpLR = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.65)
train_loss_list = []
train_acc_list = []
test_loss_list = []
test_acc_list = []

start_time = datetime.datetime.now()
print(
    f"start training，start time = {start_time}，epochs={epochs}, lr={lr}, batch_size={batch_size}, len(train_loader)={len(train_loader)}")
curr_lr = lr
for epoch in range(epochs):
    train_loss, train_acc = train_epoch(net, data_loader=train_loader, device=device, is_train=True)
    test_loss, test_acc = train_epoch(net, data_loader=test_loader, device=device, is_train=False)
    train_loss_list.append(train_loss)
    train_acc_list.append(train_acc)
    test_loss_list.append(test_loss)
    test_acc_list.append(test_acc)
    print(f"epoch:{epoch}\t train_loss:{train_loss:.10f} \t"
          f"train_acc:{train_acc} \t"
          f"epoch:{epoch}\t test_loss:{test_loss:.10f} \t"
          f"test_acc:{test_acc} \t"
          )
    # 学习率衰减策略
    # if (epoch + 1) % 2 == 0:
    #     curr_lr /= 3
    #     update_lr(optimizer, curr_lr)
    ExpLR.step()  # 学习率按照指数的形式衰减

end_time = datetime.datetime.now()
time = end_time - start_time
print(f"end training，end time = {end_time}, duration time = {time}")
loss_draw(range(1, epochs + 1), train_loss_list, "epochs", "loss", range(1, epochs + 1), test_loss_list,
          ['train', 'test'])
